/**
 * This example demonstrates how to use ROS2 to subscribe to high-level state information
 * from unitree go2 robot's /sportmodestate topic
 * Rust implementation based on the C++ version in high_level_state directory
 * 
 * Features:
 * - Subscribe to SportModeState messages
 * - Display robot position and orientation
 * - Parse and display sport state information
 * - Graceful shutdown with Ctrl+C handling
 **/

use r2r::{self, QosProfile};
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use unitree_ros2_rs::sport_client::state_machine::SportState;
use unitree_ros2_rs::utils::timespec_to_datetime;
use chrono::{DateTime, Utc};
use std::sync::Mutex;

static RUNNING: AtomicBool = AtomicBool::new(true);
static LAST_TIMESTAMP: Mutex<Option<DateTime<Utc>>> = Mutex::new(None);



/// High-level state message handler
/// Processes SportModeState messages and displays relevant information
fn high_state_handler(msg: &r2r::unitree_go::msg::SportModeState) {
    // Print timestamp information
    println!("=== SportModeState Message ===");
    println!("Timestamp: {}", timespec_to_datetime(&msg.stamp));
    println!("Raw TimeSpec: sec={}, nanosec={}", msg.stamp.sec, msg.stamp.nanosec);
    println!();
    
    // Print robot position
    println!(
        "Position: {:.3}, {:.3}, {:.3}",
        msg.position[0], msg.position[1], msg.position[2]
    );
    
    // Print robot orientation (quaternion: w, x, y, z)
    println!(
        "Quaternion: {:.3}, {:.3}, {:.3}, {:.3}",
        msg.imu_state.quaternion[0],
        msg.imu_state.quaternion[1], 
        msg.imu_state.quaternion[2],
        msg.imu_state.quaternion[3]
    );
    
    // Parse and display sport state
    if let Some(sport_state) = SportState::from_error_code(msg.error_code as i32) {
        println!("Sport State: {} ({})", sport_state.name(), msg.error_code);
    } else {
        println!("Unknown Sport State: {}", msg.error_code);
    }
    
    // Display motion mode
    let mode_name = match msg.mode {
        0 => "idle, default stand",
        1 => "balanceStand",
        2 => "pose",
        3 => "locomotion",
        4 => "reserve",
        5 => "lieDown",
        6 => "jointLock",
        7 => "damping",
        8 => "recoveryStand",
        9 => "reserve",
        10 => "sit",
        11 => "frontFlip",
        12 => "frontJump",
        13 => "frontPounce",
        _ => "unknown",
    };
    println!("Motion Mode: {} ({})", mode_name, msg.mode);
    
    // Display gait type
    let gait_name = match msg.gait_type {
        0 => "idle",
        1 => "trot",
        2 => "run",
        3 => "climb stair",
        4 => "forwardDownStair",
        9 => "adjust",
        _ => "unknown",
    };
    println!("Gait Type: {} ({})", gait_name, msg.gait_type);
    
    // Display additional information
    println!("Progress: {:.3}", msg.progress);
    println!("Body Height: {:.3}", msg.body_height);
    println!("Foot Raise Height: {:.3}", msg.foot_raise_height);
    println!(
        "Velocity: {:.3}, {:.3}, {:.3}",
        msg.velocity[0], msg.velocity[1], msg.velocity[2]
    );
    println!("Yaw Speed: {:.3}", msg.yaw_speed);
    
    // Display foot forces
    println!(
        "Foot Forces: [{}, {}, {}, {}]",
        msg.foot_force[0], msg.foot_force[1], msg.foot_force[2], msg.foot_force[3]
    );
    
    // Calculate message frequency and interval
    let current_time = chrono::Utc::now();
    let mut last_time = LAST_TIMESTAMP.lock().unwrap();
    
    if let Some(prev_time) = *last_time {
        let interval = current_time.signed_duration_since(prev_time);
        let frequency = 1000.0 / interval.num_milliseconds() as f64;
        println!("Message interval: {:.1}ms, Frequency: {:.1}Hz", 
                interval.num_milliseconds(), frequency);
    } else {
        println!("First message received");
    }
    
    *last_time = Some(current_time);
    println!("Local time: {}", current_time.format("%H:%M:%S%.3f"));
    println!("==========================================\n");
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Set up Ctrl+C handler
    ctrlc::set_handler(move || {
        println!("\nReceived Ctrl+C, shutting down safely...");
        RUNNING.store(false, Ordering::SeqCst);
    })?;

    println!("=== Go2 High Level State Monitor ===");
    println!("Subscribing to /sportmodestate topic...");
    println!("Press Ctrl+C to exit\n");

    // Create ROS2 context and node
    let ctx = r2r::Context::create()?;
    let mut node = r2r::Node::create(ctx, "go2_high_state_monitor", "")?;
    
    // Create subscriber for SportModeState
    let subscriber = node.subscribe::<r2r::unitree_go::msg::SportModeState>(
        "/sportmodestate",
        QosProfile::default(),
    )?;
    
    println!("Successfully created subscriber for /sportmodestate");
    println!("Waiting for messages...\n");
    
    // Use futures to handle async subscription
    use futures::{executor::LocalPool, stream::StreamExt, task::LocalSpawnExt};
    
    let mut pool = LocalPool::new();
    let spawner = pool.spawner();
    
    // Spawn async task to handle messages
    spawner.spawn_local(async move {
        let mut sub = subscriber;
        while let Some(msg) = sub.next().await {
            if !RUNNING.load(Ordering::SeqCst) {
                break;
            }
            high_state_handler(&msg);
        }
    }).expect("Failed to spawn subscription task");
    
    // Main loop
    while RUNNING.load(Ordering::SeqCst) {
        // Spin once to process callbacks
        node.spin_once(Duration::from_millis(10));
        
        // Run async tasks
        pool.run_until_stalled();
        
        // Small sleep to prevent excessive CPU usage
        std::thread::sleep(Duration::from_millis(1));
    }
    
    println!("High level state monitor stopped.");
    Ok(())
}